from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from base64 import b64decode
import os


class CryptoFactory:
    def __init__(self):
        self.vector_size = 16
        self.salt_size = 16
        self.tag_size = 16  # GCM tag size is 16 bytes
        self.iterations = 262144

    def derive_key(self, password: str, salt: bytes) -> bytes:
        """Derive a key from the password and salt using PBKDF2 with SHA-512."""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA512(),
            length=32,
            salt=salt,
            iterations=self.iterations,
            backend=default_backend()
        )
        key = kdf.derive(password.encode())
        return key

    def decrypt_from_bytes(self, encrypted_bytes: bytes, password: str) -> str:
        """Decrypt bytes using AES-GCM with the given password."""
        try:
            # Extract vector, salt, encrypted text, and tag
            offset = self.vector_size + self.salt_size
            vector = encrypted_bytes[:self.vector_size]
            salt = encrypted_bytes[self.vector_size:offset]
            encrypted_text_bytes = encrypted_bytes[offset:-self.tag_size]
            tag = encrypted_bytes[-self.tag_size:]

            # Derive the key using the provided password and extracted salt
            key = self.derive_key(password, salt)

            # Decrypt the text using AES-GCM
            decryptor = Cipher(
                algorithms.AES(key),
                modes.GCM(vector, tag),
                backend=default_backend()
            ).decryptor()

            decrypted_bytes = decryptor.update(encrypted_text_bytes) + decryptor.finalize()
            decrypted_text = decrypted_bytes.decode('utf-8')
            return decrypted_text
        except Exception as e:
            print("Decryption failed:", e)
            return None

    def decrypt_from_base64(self, base64_encoded: str, password: str) -> str:
        """Decode the base64-encoded encrypted text and decrypt it."""
        try:
            encrypted_bytes = b64decode(base64_encoded)
            return self.decrypt_from_bytes(encrypted_bytes, password)
        except Exception as e:
            print("Failed to decode and decrypt:", e)
            return None


# Example usage
def decrypt(content):
    crypto_factory = CryptoFactory()

    # Example encrypted base64 string and password
    # encrypted_base64 = "VhofGyUvgaCgTr//0ywq3gthiEnx5MBZXO9Twx0GVk1OBLxFnw+6U6eMDCbM2fTZmPJja3xv3Puop5oO3dc="
    password = "tan@123"

    decrypted_text = crypto_factory.decrypt_from_base64(content, password)
    return decrypted_text
